import time
import numpy as np
import gym

class SarsaAgent(object):
    def __init__(self, obs_n, act_n, learning_rate=0.01, gamma=0.9, e_greedy=0.1):
        self.act_n = act_n
        self.lr = learning_rate
        self.gamma = gamma
        self.epsilon = e_greedy
        self.Q = np.zeros((obs_n, act_n))

    def sample(self, obs):
        if np.random.uniform(0, 1) < (1.0 - self.epsilon):  # 选Q值大的动作
            action = self.predict(obs)
        else:
            action = np.random.choice(self.act_n)
        return action

    def predict(self, obs):
        Q_list = self.Q[obs, :]
        maxQ = np.max(Q_list)

        action_list = np.where(Q_list == maxQ)[0]
        action = np.random.choice(action_list)
        return action

   # 更新Q值表
    def learn(self, obs, action, reward, next_obs, next_action, done):
        predict_Q = self.Q[obs, action]
        if done:
            target_Q = reward
        else:
            target_Q = reward + self.gamma * self.Q[next_obs, next_action]
        self.Q[obs, action] = predict_Q + self.lr * (target_Q - predict_Q)


def run_episode(env, agent, render=False):
    total_steps = 0  # 总步数
    total_reward = 0  # 总奖励

    obs = env.reset()  # 初始化环境
    action = agent.sample(obs)

    while True:
        next_obs, reward, done, info = env.step(action)  # 与环境交互，并得到反馈
        next_action = agent.sample(next_obs)  # 选择下一个动作

        # 更新Q表格
        agent.learn(obs, action, reward, next_obs, next_action, done)

        action = next_action
        obs = next_obs

        total_reward += reward
        total_steps += 1

        if render:
            env.render()
        if done:
            break
    return total_reward, total_steps


def test_episode(env, agent):
    total_reward = 0
    obs = env.reset()

    while True:
        action = agent.predict(obs)  # greedy
        next_obs, reward, done, info = env.step(action)
        total_reward += reward
        obs = next_obs
        time.sleep(0.5)
        if done:
            break
            return total_reward

env = gym.make("CliffWalking-v0")  # 引入环境

agent = SarsaAgent(
        obs_n=env.observation_space.n,
        act_n=env.action_space.n,
        learning_rate=0.1,
        gamma=0.9,
        e_greedy=0.1)

is_render = False
# 迭代500
for episode in range(500):
       ep_reward, ep_steps = run_episode(env, agent,  is_render)
       print('Episode %s: steps = %s, reward = %.lf' % (episode, ep_steps, ep_reward))

test_reward = test_episode(env, agent)
print('test reward = %.1f' % (test_reward))

